import os

from langchain_community.embeddings import DashScopeEmbeddings
from langchain_chroma import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_openai import ChatOpenAI

os.environ['LANGCHAIN_TRACING_V2'] = 'true'
os.environ['LANGCHAIN_PROJECT'] = 'LLMDEMO'
os.environ['LANGCHAIN_API_KEY'] = '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'
# 初始化带模板支持的LLM
# pip install redis
llm = ChatOpenAI(
    api_key="sk-a3f7718fb81f43b2915f0a6483b6661b",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    model="qwen-plus",  # 此处以qwen-plus为例，您可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
    # other params...
)
#pip install pypdf
# pip install langchain-chroma
# pip install chroma-hnswlib==0.7.5 chromadb==0.5.4 版本不兼容报错问题
# pip install dashscope
#pip install langchain-chroma

# RAG流程 加载文档(PDF、CSV等) 向量化文档 相似性检索
# loader = PyPDFLoader(r"D:\pythonWork\python_demo\pdf\问答题.pdf")
# pages = loader.load_and_split()
# text_spliter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=100, length_function=len,add_start_index=True)
# doc = text_spliter.create_documents([pages[0].page_content])
# for doc in doc:
#     print(doc.page_content)
#     print("-----------------")
#加载文档并把向量数据库保存到本地，保存在embedding文件夹下
db = Chroma(embedding_function=DashScopeEmbeddings(),persist_directory="embedding")
retriever = RunnableLambda(db.similarity_search).bind(k=1)
print(retriever)
prompt_template = """
你是一个智能问答机器人
你的任务是根据下述给定的已知信息回答用户问题。
确保你的回复完全依据一下已知信息，不要编造答案。
如果下述已知信息不足以回答用户的回复，请直接回复"我无法回答你的问题"
已知信息：
{info}
用户问题：
{input}
请用中文回答。
"""
prompt_template = ChatPromptTemplate.from_messages([("human",prompt_template)])
# 在chain中，runnables模块中的RunnablePassthrough可用于格式化上一个输出，从而匹配下一个输入的需要的格式
chain = {"input": RunnablePassthrough(), "info": retriever} | prompt_template | llm
while True:
    text = input("请输入问题：")
    if text == "exit":
        break
    print(chain.invoke(text).content)
